import json
from typing import List
from os.path import abspath, join
import re
from backend.experiment.clue.chid.scheme0.implement import CHIDImplement
from backend.experiment.tokenizer import tokenizer
from backend.experiment.model import vocab_name
from backend.experiment.framework.shots_prompt.implement import ShotMode
from backend.experiment.framework.task import PerplexityChoiceMakingTask

test_dev = True
insert_start_token = True
only_calculate_last = True
seed = 42
task_handler = CHIDImplement.get_task_handler()


class CHIDSample:
    def __init__(
            self, text: str, candidates: List[str], answer: str, task
    ):
        self.__candidates = candidates
        self.__answer = answer
        self.__text = text
        self.__task = task

    @property
    def candidates(self) -> List[str]:
        return self.__candidates

    @property
    def answer(self) -> str:
        return self.__answer

    @property
    def text(self) -> str:
        return self.__text

    @property
    def task(self) -> PerplexityChoiceMakingTask:
        return self.__task


def initialize() -> List[CHIDSample]:
    with open(abspath(
            join(__file__, '../../experiment/clue/chid/dev.json')),
            'r', encoding='utf8') as file:
        records = [json.loads(each) for each in file.readlines()]

    with open(abspath(
            join(__file__, '../../experiment/clue/chid/dev_answer.json')),
            'r', encoding='utf8') as file:
        ans_records: dict = json.loads(file.read())

    cache_filename = f'{vocab_name}_zero_shot_{seed}'
    if insert_start_token:
        cache_filename += '_start_token'

    if only_calculate_last:
        cache_filename += '_only_last_shot'

    test_tasks = CHIDImplement(
        tokenizer=tokenizer,
        test_dev=test_dev,
        seed=seed
    ).get_tasks(
        mode=ShotMode.FEW_SHOT,
        cache_filename=cache_filename + '.json',
        insert_start_token=insert_start_token,
        only_calculate_last=only_calculate_last
    )

    _samples = []
    i = 0
    for each in records:
        candidates = each['candidates']
        for content in each['content']:
            found = re.findall('#idiom[0-9]*#', content)
            assert len(found) > 0
            if len(found) > 1:
                # 跳过有多个空的样例
                for _ in found:
                    i += 1
            else:
                if test_tasks[i] is not None:
                    text = re.sub('#idiom[0-9]*#', '[MASK]', content)
                    _samples.append(
                        CHIDSample(
                            candidates=candidates,
                            answer=candidates[int(ans_records[found[0]])],
                            text=text,
                            task=test_tasks[i]
                        )
                    )
                i = i + 1
    return _samples


samples = initialize()
